#!/usr/bin/env python3
#
# @file OnnxInfer.py
# @brief
# @author QRS
# @version 1.0
# @date 2024-04-03 13:28

import os
import time
import cv2
import numpy as np
import onnxruntime as ort

from Utils import Logger

onnx_file_path = os.environ.get('ONNX_FILE_PATH', 'welding_spots_seg.onnx')

mean = np.array((0, 0, 0)).reshape((3, 1, 1))
std = np.array((1, 1, 1)).reshape((3, 1, 1))
input_size = (224, 224)

session = ort.InferenceSession(onnx_file_path)


def model_infer(images, debug=False):
    if debug:
        t0 = time.time()
    inputs = []
    for img in images:
        img = cv2.resize(img, input_size)
        img = np.transpose(img, [2, 0, 1]) / 255.0
        img = (img - mean) / std
        inputs.append(np.array(img).astype('float32'))

    input_dict = {session.get_inputs()[0].name: np.stack(inputs, axis=0)}
    if debug:
        t1 = time.time()
    outputs = session.run(None, input_dict)[0]
    if debug:
        t2 = time.time()
        Logger.info(f'preprocess time {t1 - t0}, inference time {t2 - t1}, all time {t2 - t0}')

    preds = outputs[:, 0, :, :]
    preds[np.nonzero(preds < 0.5)] = 0.0
    preds[np.nonzero(preds >= 0.5)] = 255.
    return preds.astype("uint8")
